# 人工智能NLP-Agent数字人项目-04-基金数据问答任务工单V1.1-2.14
import sqlite3
import pandas as pd
import utils.configFinRAG as configFinRAG
from FinSQL_01_generate import generate_sql
from FinSQL_02_query import query_db
from FinSQL_03_answer_from_SQL import generate_answer
from utils.instances import TOKENIZER, LLM

# 定义全局变量
g_example_question_list = []
g_example_sql_list = []
g_example_fa_list = []
g_example_info_list = []
g_example_token_list = []

def load_sql_examples():
    """
    加载SQL示例数据，包括问题、SQL语句、资料、FA和token列表
    """
    global g_example_question_list, g_example_sql_list, g_example_fa_list, g_example_info_list, g_example_token_list
    try:
        if len(g_example_question_list) <= 0:
            sql_examples_file = pd.read_csv(configFinRAG.sql_examples_path, delimiter=",", header=0)
            for _, row in sql_examples_file.iterrows():
                g_example_question_list.append(row['问题'])
                g_example_sql_list.append(row['SQL'])
                g_example_info_list.append(row['资料'])
                g_example_fa_list.append(row['FA'])
                tokens = TOKENIZER(row['问题'])
                g_example_token_list.append(tokens['input_ids'])
    except FileNotFoundError:
        print(f"未找到文件: {configFinRAG.sql_examples_path}")
    except Exception as e:
        print(f"加载SQL示例数据时发生错误: {e}")

def sql_retrieve_chain(query):
    """
    SQL检索链，包括生成SQL、查询数据库和生成答案
    """
    # 加载SQL示例数据
    load_sql_examples()

    # 生成SQL
    try:
        result_prompt, sql = generate_sql(query, LLM, g_example_question_list, g_example_sql_list, g_example_token_list)
    except Exception as e:
        print(f"生成SQL时发生错误: {e}")
        return None

    # 连接数据库并执行查询
    try:
        db_path = r'C:\Users\langzi\Desktop\local_human_number\数字人\Fay-fay-agent-edition0830\agent\tools\data\博金杯比赛数据.db'
        with sqlite3.connect(db_path) as conn:
            cs = conn.cursor()
            success_flag, exc_result = query_db(sql, cs)
    except sqlite3.Error as e:
        print(f"查询数据库时发生错误: {e}")
        return None

    # 生成答案
    try:
        answer = generate_answer(query, exc_result, LLM, g_example_question_list, g_example_info_list, g_example_fa_list,
                                 g_example_token_list)
        return answer
    except Exception as e:
        print(f"生成答案时发生错误: {e}")
        return None